//Библиотека Qt.
#include <QCoreApplication>
#include <QDate>
#include <QDir>


//Внутренняя библиотека.
#include "filewriter.h"
#include "viddevice.h"
#include "vidframe.h"


//Директивы препроцессора.
#define AVSTREAM_CODEC_PTR(avs_ptr) (avs_ptr->codec)


//Дополнительная информация, которую необходимо указывать для mpeg1 видеофайлов.
static unsigned char mpeg1_trailer[] = {0x00, 0x00, 0x01, 0xb7};


//===========================================================================//
//=Заголовок: static int mpeg1WriteTrailer()
//=Параметры:
//=Описание: Функция записи дополнительной информации в mpeg1 видеофайл.
//===========================================================================//
static int mpeg1WriteTrailer(AVFormatContext *s)
{
put_buffer(s->pb, mpeg1_trailer, 4);
put_flush_packet(s->pb);
return 0;
}


//===========================================================================//
//=Заголовок: FileWriter::FileWriter() - public
//=Параметры:
//=Описание: Конструктор.
//===========================================================================//
FileWriter::FileWriter(QString devname, VidDevice *parent) : QThread(parent)
{
moveToThread(this);

deviceName = devname;

isRecording = false;

fmtCtx = NULL; videoStream = NULL; codecCtx = NULL; vidCodec = FLV;
outBuf = NULL; outBufSize = 0;

vidFrame = NULL;

//Инициализировать FFMPEG.
av_register_all();
}


//===========================================================================//
//=Заголовок: FileWriter::~FileWriter() - public
//=Параметры:
//=Описание: Деструктор.
//===========================================================================//
FileWriter::~FileWriter() {if(isRecording) stop(); wait();}


//===========================================================================//
//=Заголовок: void FileWriter::setDeviceName() - public
//=Параметры:
//=Описание: Функция установки наименования устройства.
//===========================================================================//
void FileWriter::setDeviceName(QString devname)
{mutex.lock(); deviceName = devname; mutex.unlock();}


//===========================================================================//
//=Заголовок: void FileWriter::write() - public
//=Параметры:
//=Описание: Функция записи данных в файл.
//===========================================================================//
void FileWriter::write(VidFrame *vidframe)
{
if(!vidframe) return;
vidFrame = vidframe;
QThread::start();
}


//===========================================================================//
//=Заголовок: void FileWriter::stopWrite() - public
//=Параметры:
//=Описание: Функция остановки записи данных в файл.
//===========================================================================//
void FileWriter::stopWrite() {wait(); vidFrame = NULL; QThread::start();}


//===========================================================================//
//=Заголовок: void FileWriter::run() - protected
//=Параметры:
//=Описание: Функция активации нити.
//===========================================================================//
void FileWriter::run()
{
if(vidFrame && !isRecording) return start(vidFrame);
if(!vidFrame && isRecording) return stop();
if(!isRecording) return;

AVFrame *srcpic = avcodec_alloc_frame();
if(!srcpic) return;
avpicture_fill((AVPicture*)srcpic, vidFrame->getData(), vidFrame->getFFMPEGPixFormat(), vidFrame->getWidth(), vidFrame->getHeight());

AVFrame *dstpic = avcodec_alloc_frame();
if(!dstpic) return;
avpicture_fill((AVPicture*)dstpic, outBuf, codecCtx->pix_fmt, codecCtx->width, codecCtx->height);

int ret;

int out_size = avcodec_encode_video(codecCtx, outBuf, outBufSize, srcpic);
if(out_size != 0)
	{
	AVPacket pkt;
	av_init_packet(&pkt);
	pkt.stream_index = videoStream->index;
	if(codecCtx->coded_frame->pts != (unsigned)(0x8000000000000000LL))
		pkt.pts = av_rescale_q(codecCtx->coded_frame->pts, codecCtx->time_base, videoStream->time_base);
	if(codecCtx->coded_frame->key_frame)
		pkt.flags |= PKT_FLAG_KEY;

	pkt.data = outBuf;
	pkt.size = out_size;
	ret = av_interleaved_write_frame(fmtCtx, &pkt);
}

av_free(srcpic);
av_free(dstpic);
}


//===========================================================================//
//=Заголовок: void FileWriter::start() - private
//=Параметры:
//=Описание: Функция активации записи в файл.
//===========================================================================//
void FileWriter::start(VidFrame *vidframe)
{
//Bitrate to be used by the ffmpeg encoder (default: 400000)
//This option is ignored if ffmpeg_variable_bitrate is not 0 (disabled)
int bps = 2500000;

//Enables and defines variable bitrate for the ffmpeg encoder.
//ffmpeg_bps is ignored if variable bitrate is enabled.
//Valid values: 0 (default) = fixed bitrate defined by ffmpeg_bps,
//or the range 2 - 31 where 2 means best quality and 31 is worst.
int vbr = 0;

//Maximum number of frames to be captured per second.
//Valid range: 2-30. Default: 20 (almost no limit).
int fps = 25;

fmtCtx = (AVFormatContext *)av_mallocz(sizeof(AVFormatContext));
if(!fmtCtx) return;

//Определить путь к архиву видеозаписей.
QString fname = QCoreApplication::applicationDirPath() + "/" + deviceName + "/" + QDate::currentDate().toString("dd.MM.yyyy");

//Создать путь к архиву видеозаписей.
QDir dir; dir.mkpath(fname);

//Определить имя файла для записи.
fname += "/" + deviceName + "-" + QTime::currentTime().toString("hh:mm:ss");

fmtCtx->oformat = getOutputFormat(&fname);
if(!fmtCtx->oformat) {av_free(fmtCtx); fmtCtx = NULL; return;}

snprintf(fmtCtx->filename, sizeof(fmtCtx->filename), "%s", (const char*)fname.toAscii());

if(fmtCtx->oformat->video_codec != CODEC_ID_NONE)
	{
	videoStream = av_new_stream(fmtCtx, 0);
	if(!videoStream) return;
}
else return;

AVCodecContext *codecctx;
codecCtx = codecctx = AVSTREAM_CODEC_PTR(videoStream);
codecctx->codec_id = fmtCtx->oformat->video_codec;
codecctx->codec_type = CODEC_TYPE_VIDEO;
bool is_mpeg1 = codecctx->codec_id == CODEC_ID_MPEG1VIDEO;

if(vidCodec == FFV1) codecctx->strict_std_compliance = -2;

codecctx->bit_rate = bps;
codecctx->width = vidframe->getWidth();
codecctx->height = vidframe->getHeight();

codecctx->time_base.num = 1;
codecctx->time_base.den = fps;

if(vbr) codecctx->flags |= CODEC_FLAG_QSCALE;

codecctx->gop_size = 10;//is_mpeg1 ? 10 : 12;

if(fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
	codecctx->flags |= CODEC_FLAG_GLOBAL_HEADER;

if(av_set_parameters(fmtCtx, NULL) < 0) return;

AVCodec *codec = avcodec_find_encoder(codecctx->codec_id);
if(!codec) return;

codecctx->pix_fmt = vidframe->getFFMPEGPixFormat();

if(avcodec_open(codecctx, codec) < 0) return;

if(!(fmtCtx->oformat->flags & AVFMT_RAWPICTURE))
	{
	outBufSize = avpicture_get_size(vidframe->getFFMPEGPixFormat(), vidframe->getWidth(), vidframe->getHeight());
	outBuf = (uint8_t *)malloc(outBufSize);
	if(!outBuf) return;
}

if(!(fmtCtx->oformat->flags & AVFMT_NOFILE))
	{
	char file_proto[256];
	if(is_mpeg1) snprintf(file_proto, sizeof(file_proto), "v:%s", (const char*)fname.toAscii());
	else snprintf(file_proto, sizeof(file_proto), "%s", (const char*)fname.toAscii());
	if(url_fopen(&fmtCtx->pb, file_proto, URL_WRONLY) < 0) return;
}

av_write_header(fmtCtx);

isRecording = true;
}


//===========================================================================//
//=Заголовок: void FileWriter::stop() - private
//=Параметры:
//=Описание: Функция остановки записи в файл.
//===========================================================================//
void FileWriter::stop()
{
if(!isRecording) return;

if(videoStream)
	{
	avcodec_close(AVSTREAM_CODEC_PTR(videoStream));
	free(outBuf); outBuf = NULL;
}
videoStream = NULL;

av_write_trailer(fmtCtx);

for(unsigned int i = 0; i < fmtCtx->nb_streams; i++)
	av_freep(&fmtCtx->streams[i]);

if(!(fmtCtx->oformat->flags & AVFMT_NOFILE)) url_fclose(fmtCtx->pb);

if(fmtCtx) av_free(fmtCtx); fmtCtx = NULL;
if(codecCtx) av_free(codecCtx); codecCtx = NULL;

isRecording = false;
}


//===========================================================================//
//=Заголовок: AVOutputFormat *FileWriter::initOutputFormat() - private
//=Параметры: *path_fbasename - путь и базовое (без точки и расширения)
//=                             имя файла для записи.
//=Описание: Функция возврата выходного видеоформата.
//===========================================================================//
AVOutputFormat *FileWriter::getOutputFormat(QString *path_fbasename)
{
AVOutputFormat *outputfmt = NULL;

QString file_ext;

switch(vidCodec)
	{
	case MPEG1:
		file_ext = ".mpg";
		outputfmt = guess_format("mpeg1video", NULL, NULL);
		if(outputfmt) outputfmt->write_trailer = mpeg1WriteTrailer;
	break;

	case MPEG4:
		file_ext = ".avi";
		outputfmt = guess_format("avi", NULL, NULL);
	break;

	case MSMPEG4:
		file_ext = ".avi";
		outputfmt = guess_format("avi", NULL, NULL);
		if(outputfmt) outputfmt->video_codec = CODEC_ID_MSMPEG4V2;
	break;

	case SWF:
		file_ext = ".swf";
		outputfmt = guess_format("swf", NULL, NULL);
	break;

	case FLV:
		file_ext = ".flv";
		outputfmt = guess_format("flv", NULL, NULL);
	break;

	case FFV1:
		file_ext = ".avi";
		outputfmt = guess_format("avi", NULL, NULL);
		if(outputfmt) outputfmt->video_codec = CODEC_ID_FLV1;
	break;

	case MOV:
		file_ext = ".mov";
		outputfmt = guess_format("mov", NULL, NULL);
	break;
}

(*path_fbasename) += file_ext;

return outputfmt;
}
